#include <algorithm>
#include <memory>
#include "catalog/column.h"
#include "catalog/schema.h"
#include "common/exception.h"
#include "common/macros.h"
#include "execution/expressions/column_value_expression.h"
#include "execution/expressions/comparison_expression.h"
#include "execution/expressions/constant_value_expression.h"
#include "execution/expressions/logic_expression.h"
#include "execution/plans/abstract_plan.h"
#include "execution/plans/filter_plan.h"
#include "execution/plans/hash_join_plan.h"
#include "execution/plans/nested_loop_join_plan.h"
#include "execution/plans/projection_plan.h"
#include "optimizer/optimizer.h"
#include "type/type_id.h"

namespace bustub {

auto Optimizer::OptimizeNLJAsHashJoin(const AbstractPlanNodeRef &plan) -> AbstractPlanNodeRef {
  // TODO(student): implement NestedLoopJoin -> HashJoin optimizer rule
  // Note for 2023 Spring: You should at least support join keys of the form:
  // 1. <column expr> = <column expr>
  // 2. <column expr> = <column expr> AND <column expr> = <column expr>
  auto is_col_equal_expr = [](const AbstractExpression *base_expr) -> bool {
    const auto *comp_expr = dynamic_cast<const ComparisonExpression *>(base_expr);
    if (comp_expr != nullptr && comp_expr->comp_type_ == ComparisonType::Equal) {
      const auto &childs = comp_expr->GetChildren();
      if (const auto *lch_expr = dynamic_cast<const ColumnValueExpression *>(childs[0].get()); lch_expr != nullptr) {
        if (const auto *rch_expr = dynamic_cast<const ColumnValueExpression *>(childs[1].get()); rch_expr != nullptr) {
          return true;
        }
      }
    }
    return false;
  };

  std::vector<AbstractPlanNodeRef> children;
  for (const auto &child : plan->GetChildren()) {
    children.emplace_back(OptimizeNLJAsHashJoin(child));
  }
  auto optimized_plan = plan->CloneWithChildren(std::move(children));

  if (optimized_plan->GetType() == PlanType::NestedLoopJoin) {
    const auto &nlj_plan = dynamic_cast<const NestedLoopJoinPlanNode &>(*optimized_plan);
    // Has exactly two children
    BUSTUB_ENSURE(nlj_plan.children_.size() == 2, "NLJ should have exactly 2 children.");
    const AbstractExpression *expr = nlj_plan.Predicate().get();
    if (expr == nullptr) {
      return optimized_plan;
    }
    if (is_col_equal_expr(expr)) {
      const auto *lch_expr = reinterpret_cast<const ColumnValueExpression *>(expr->GetChildAt(0).get());
      const auto *rch_expr = reinterpret_cast<const ColumnValueExpression *>(expr->GetChildAt(1).get());
      std::vector<AbstractExpressionRef> l_exprs;
      std::vector<AbstractExpressionRef> r_exprs;
      if (lch_expr->GetTupleIdx() == 0 && rch_expr->GetTupleIdx() == 1) {
        l_exprs.push_back(expr->GetChildAt(0));
        r_exprs.push_back(expr->GetChildAt(1));
      } else if (lch_expr->GetTupleIdx() == 1 && rch_expr->GetTupleIdx() == 0) {
        l_exprs.push_back(expr->GetChildAt(1));
        r_exprs.push_back(expr->GetChildAt(0));
      }
      return std::make_shared<HashJoinPlanNode>(nlj_plan.output_schema_, nlj_plan.GetLeftPlan(),
                                                nlj_plan.GetRightPlan(), l_exprs, r_exprs, nlj_plan.GetJoinType());

    } else if (const auto *logic_expr = dynamic_cast<const LogicExpression *>(expr); logic_expr != nullptr) {
      if (logic_expr->logic_type_ == LogicType::And) {
        const auto &childs = expr->GetChildren();
        if (childs.size() == 2 && is_col_equal_expr(childs[0].get()) && is_col_equal_expr(childs[1].get())) {
          std::vector<AbstractExpressionRef> l_exprs;
          std::vector<AbstractExpressionRef> r_exprs;
          for (const auto &ch : childs) {
            const auto *equal_expr = ch.get();
            const auto *lch_expr = reinterpret_cast<const ColumnValueExpression *>(equal_expr->GetChildAt(0).get());
            const auto *rch_expr = reinterpret_cast<const ColumnValueExpression *>(equal_expr->GetChildAt(1).get());
            if (lch_expr->GetTupleIdx() == 0 && rch_expr->GetTupleIdx() == 1) {
              l_exprs.push_back(equal_expr->GetChildAt(0));
              r_exprs.push_back(equal_expr->GetChildAt(1));
            } else if (lch_expr->GetTupleIdx() == 1 && rch_expr->GetTupleIdx() == 0) {
              l_exprs.push_back(equal_expr->GetChildAt(1));
              r_exprs.push_back(equal_expr->GetChildAt(0));
            }
          }
          return std::make_shared<HashJoinPlanNode>(nlj_plan.output_schema_, nlj_plan.GetLeftPlan(),
                                                    nlj_plan.GetRightPlan(), l_exprs, r_exprs, nlj_plan.GetJoinType());
        }
      }
    }
  }

  return optimized_plan;
}

}  // namespace bustub
